from sklearn.linear_model import Ridge
import h5py
import pandas as pd
from matplotlib import pyplot as plt
import os

def process_dir(path, dest_path, y_col, get_mask):
    os.makedirs(dest_path, exist_ok=True)

    for fname in os.listdir(path):
        p = os.path.join(path, fname)
        f = h5py.File(p, 'r')

        ts = f['ts'][...]
        x = f['x'][...]
        y = f[y_col][...]
        mask = get_mask(f)

        ts = ts[mask]
        x = x[mask]
        y = y[mask]
        if len(ts)==0:
          continue

        a = Ridge(alpha=0.1)
        a.fit(x, y)

        p = os.path.join(dest_path, fname)
        print('writing ', p)
        f = h5py.File(p, 'w')
        f['weights'] = a.coef_
        f['bias'] = a.intercept_


def get_mask1(f):
  return (f['y1'][...]>=0) & (f['y2'][...]>=0)

def get_mask2(f):
  return (f['y'][...]>=0)

process_dir('../coin_binary/latency-dataset/order-submit', '../coin_binary/latency-model/order-submit-to-accepted', 'y1', get_mask1)
process_dir('../coin_binary/latency-dataset/order-submit', '../coin_binary/latency-model/order-submit-to-first-fill', 'y2', get_mask1)
process_dir('../coin_binary/latency-dataset/order-cancel', '../coin_binary/latency-model/cancel-submit-to-confirmed', 'y', get_mask2)
